#include <linux/kprobes.h>
#include <linux/init.h>
#include <linux/module.h>
#include <linux/slab.h>
#include <linux/proc_fs.h>
#include <linux/uaccess.h>
#include <linux/list.h>
#include <linux/spinlock.h>

struct malloc_entry
{
    unsigned long address;
    size_t size;
    struct list_head list;
};

static LIST_HEAD(malloc_list);
static DEFINE_SPINLOCK(malloc_lock);

// kmalloc pre_handler
static int kmalloc_pre_handler(struct kprobe *p, struct pt_regs *regs)
{
    size_t size = regs->si;           // Size argument
    unsigned long address = regs->ax; // Return value (address)

    struct malloc_entry *entry = kmalloc(sizeof(*entry), GFP_ATOMIC);
    if (!entry)
    {
        printk(KERN_ERR "Failed to allocate malloc_entry\n");
        return 0;
    }

    entry->address = address;
    entry->size = size;

    spin_lock(&malloc_lock);
    list_add_tail(&entry->list, &malloc_list);
    spin_unlock(&malloc_lock);

    printk(KERN_INFO "kmalloc: address=0x%lx, size=%zu\n", address, size);
    return 0;
}

// kfree pre_handler
static int kfree_pre_handler(struct kprobe *p, struct pt_regs *regs)
{
    unsigned long address = regs->di; // Address argument

    struct malloc_entry *entry, *tmp;
    bool found = false;

    spin_lock(&malloc_lock);
    list_for_each_entry_safe(entry, tmp, &malloc_list, list)
    {
        if (entry->address == address)
        {
            list_del(&entry->list);
            kfree(entry);
            found = true;
            break;
        }
    }
    spin_unlock(&malloc_lock);

    if (!found)
    {
        printk(KERN_ERR "kfree: address=0x%lx not found in malloc records\n", address);
    }
    else
    {
        printk(KERN_INFO "kfree: address=0x%lx\n", address);
    }
    return 0;
}

static struct kprobe kp_km = {
    .symbol_name = "__kmalloc",
    .pre_handler = kmalloc_pre_handler,
};

static struct kprobe kp_kf = {
    .symbol_name = "kfree",
    .pre_handler = kfree_pre_handler,
};

static int __init kprobe_init(void)
{
    register_kprobe(&kp_km);
    register_kprobe(&kp_kf);
    printk(KERN_INFO "Kprobes registered\n");
    return 0;
}

static void __exit kprobe_exit(void)
{
    unregister_kprobe(&kp_km);
    unregister_kprobe(&kp_kf);
    printk(KERN_INFO "Kprobes unregistered\n");
}

module_init(kprobe_init);
module_exit(kprobe_exit);

MODULE_LICENSE("GPL");
MODULE_DESCRIPTION("Kprobe example to monitor kmalloc and kfree");
MODULE_AUTHOR("Your Name");